package cn.edu.hitsz.compiler.asm;

import cn.edu.hitsz.compiler.NotImplementedException;
import cn.edu.hitsz.compiler.ir.*;
import cn.edu.hitsz.compiler.utils.FileUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;


/**
 * TODO: 实验四: 实现汇编生成
 * <br>
 * 在编译器的整体框架中, 代码生成可以称作后端, 而前面的所有工作都可称为前端.
 * <br>
 * 在前端完成的所有工作中, 都是与目标平台无关的, 而后端的工作为将前端生成的目标平台无关信息
 * 根据目标平台生成汇编代码. 前后端的分离有利于实现编译器面向不同平台生成汇编代码. 由于前后
 * 端分离的原因, 有可能前端生成的中间代码并不符合目标平台的汇编代码特点. 具体到本项目你可以
 * 尝试加入一个方法将中间代码调整为更接近 risc-v 汇编的形式, 这样会有利于汇编代码的生成.
 * <br>
 * 为保证实现上的自由, 框架中并未对后端提供基建, 在具体实现时可自行设计相关数据结构.
 *
 * @see AssemblyGenerator#run() 代码生成与寄存器分配
 */
public class AssemblyGenerator {

    HashMap<IRValue, Integer> lastUse = new HashMap<>(); // 变量最后一次使用的位置(instructionList列表中的下标)
    List<Instruction> instructionList = new ArrayList<>(); // 预处理后的指令列表
    HashMap<IRVariable, String> regMap = new HashMap<>(); // 各个变量占用寄存器的情况
    List<String> freeRegs = new ArrayList<>(List.of("t0", "t1", "t2", "t3", "t4", "t5", "t6")); // 可用的寄存器
    List<String> asmList = new ArrayList<>(); // 生成的汇编列表

    /**
     * 加载前端提供的中间代码
     * <br>
     * 视具体实现而定, 在加载中或加载后会生成一些在代码生成中会用到的信息. 如变量的引用
     * 信息. 这些信息可以通过简单的映射维护, 或者自行增加记录信息的数据结构.
     *
     * @param originInstructions 前端提供的中间代码
     */
    public void loadIR(List<Instruction> originInstructions) {
        // TODO: 读入前端提供的中间代码并生成所需要的信息
        // 记录各个变量的最后一次使用时间
        int cnt = 0;
        for (Instruction inst : originInstructions) {
            switch (inst.getKind()) {
                case ADD -> {
                    /**
                     * - 左边是立即数
                     * --- 右边是立即数 -> 转为MOV
                     * --- 右边是变量   -> 交换位置, 记录变量
                     * - 左边是变量
                     * --- 右边是变量  -> 记录两个变量最后有效时间
                     * --- 右边是立即数 -> 记录左边变量最后有效时间 (生成addi的工作交给run函数)
                     */
                    if (inst.getLHS() instanceof IRImmediate li) {
                        if (inst.getRHS() instanceof IRImmediate ri) {
                            instructionList.add(Instruction.createMov(inst.getResult(),
                                    IRImmediate.of(li.getValue() + ri.getValue())));
                        } else {
                            // 进行交换
                            lastUse.put(inst.getRHS(), cnt);
                            instructionList.add(Instruction.createAdd(inst.getResult(), inst.getRHS(), inst.getLHS()));
                        }
                    } else {
                        lastUse.put(inst.getLHS(), cnt);
                        if (inst.getRHS().isIRVariable()) {
                            lastUse.put(inst.getRHS(), cnt);
                        }
                        instructionList.add(inst);
                    }
                    lastUse.put(inst.getResult(), cnt);
                }
                case SUB -> {
                    /**
                     * - 左边是立即数
                     * --- 右边是立即数 -> 转为MOV
                     * --- 右边是变量   -> 重新生成MOV和SUB两条指令
                     * - 左边是变量
                     * --- 右边是变量  -> 记录两个变量最后有效时间
                     * --- 右边是立即数 -> 记录左边变量最后有效时间 (生成addi的工作交给run函数)
                     */
                    if (inst.getLHS() instanceof IRImmediate li) {
                        if (inst.getRHS() instanceof IRImmediate ri) {
                            instructionList.add(Instruction.createMov(inst.getResult(),
                                    IRImmediate.of(li.getValue() - ri.getValue())));
                        } else {
                            // 重新生成MOV和SUB两条指令
                            IRVariable tmp = IRVariable.temp();
                            instructionList.add(Instruction.createMov(tmp, inst.getLHS()));
                            cnt++;
                            instructionList.add(Instruction.createSub(inst.getResult(), tmp, inst.getRHS()));
                            lastUse.put(tmp, cnt);
                            lastUse.put(inst.getRHS(), cnt);
                        }
                    } else {
                        lastUse.put(inst.getLHS(), cnt);
                        if (inst.getRHS().isIRVariable()) {
                            lastUse.put(inst.getRHS(), cnt);
                        }
                        instructionList.add(inst);
                    }
                    lastUse.put(inst.getResult(), cnt);
                }
                case MUL -> {
                    /**
                     * - 左边是立即数
                     * --- 右边是立即数 -> 转为MOV
                     * --- 右边是变量   -> 重新生成MOV和MUL两条指令
                     * - 左边是变量
                     * --- 右边是变量  -> 记录两个变量最后有效时间
                     * --- 右边是立即数 -> 重新生成MOV和MUL两条指令
                     */
                    if (inst.getLHS() instanceof IRImmediate li) {
                        if (inst.getRHS() instanceof IRImmediate ri) {
                            instructionList.add(Instruction.createMov(inst.getResult(),
                                    IRImmediate.of(li.getValue() * ri.getValue())));
                        } else {
                            // 重新生成MOV和SUB两条指令
                            IRVariable tmp = IRVariable.temp();
                            instructionList.add(Instruction.createMov(tmp, inst.getLHS()));
                            cnt++;
                            instructionList.add(Instruction.createMul(inst.getResult(), tmp, inst.getRHS()));
                            lastUse.put(tmp, cnt);
                            lastUse.put(inst.getRHS(), cnt);
                        }
                    } else {
                        lastUse.put(inst.getLHS(), cnt);
                        if (inst.getRHS().isIRVariable()) {
                            lastUse.put(inst.getRHS(), cnt);
                            instructionList.add(inst);
                        } else {
                            // 重新生成MOV和SUB两条指令
                            IRVariable tmp = IRVariable.temp();
                            instructionList.add(Instruction.createMov(tmp, inst.getRHS()));
                            cnt++;
                            instructionList.add(Instruction.createMul(inst.getResult(), inst.getLHS(), tmp));
                            lastUse.put(tmp, cnt);
                            lastUse.put(inst.getLHS(), cnt);
                        }
                    }
                    lastUse.put(inst.getResult(), cnt);
                }
                case MOV -> {
                    // 只需记录变量使用时间
                    if (inst.getFrom().isIRVariable()) {
                        lastUse.put(inst.getFrom(), cnt);
                    }
                    lastUse.put(inst.getResult(), cnt);
                    instructionList.add(inst);
                }
                case RET -> {
                    // 只需记录变量使用时间
                    if (inst.getReturnValue().isIRVariable()) {
                        lastUse.put(inst.getReturnValue(), cnt);
                    }
                    instructionList.add(inst);
//                    FileUtils.writeLines("data/out/testAsm.txt", instructionList.stream().map(Instruction::toString).toList());
//                    System.out.println(lastUse);
                    return;
                }
            }
            cnt++;
        }
    }

    /**
     * 获取变量所属的寄存器 会尝试分配寄存器
     *
     * @param variable IRVariable
     * @param cnt      当前第几条指令(0开始)
     * @return 寄存器名
     */
    String allocReg(IRVariable variable, int cnt) {
        String reg;
        if (regMap.containsKey(variable)) {
            // 已分配寄存器 需要检查是否是最后一次使用, 如果是的话释放该寄存器
            reg = regMap.get(variable);
            if (cnt == lastUse.get(variable)) {
                freeRegs.add(reg);
            }
            return regMap.get(variable);
        }
        // 不完备的寄存器分配
        if (freeRegs.size() == 0) {
            throw new RuntimeException("Not Enough regs" + cnt);
        }
        // 分配空闲寄存器 上面已经检测了变量最后一次使用 这里无需考虑
        reg = freeRegs.remove(0);
        regMap.put(variable, reg);
        return reg;
    }

    /**
     * 执行代码生成.
     * <br>
     * 根据理论课的做法, 在代码生成时同时完成寄存器分配的工作. 若你觉得这样的做法不好,
     * 也可以将寄存器分配和代码生成分开进行.
     * <br>
     * 提示: 寄存器分配中需要的信息较多, 关于全局的与代码生成过程无关的信息建议在代码生
     * 成前完成建立, 与代码生成的过程相关的信息可自行设计数据结构进行记录并动态维护.
     */
    public void run() {
        // TODO: 执行寄存器分配与代码生成
        asmList.add(".text");
        for (int i = 0; i < instructionList.size(); i++) {
            Instruction inst = instructionList.get(i);
            switch (inst.getKind()) {
                case ADD, SUB, MUL -> {
                    String instName = inst.getKind().name().toLowerCase();
                    // 先检查lhs和rhs的寄存器分配情况, 如果其为最后一次使用, res可以重复利用
                    String rhs;
                    if (inst.getRHS() instanceof IRImmediate imm) {
                        rhs = imm.toString();
                        instName += "i";
                    } else {
                        rhs = allocReg((IRVariable) inst.getRHS(), i);
                    }
                    String lhs = allocReg((IRVariable) inst.getLHS(), i);
                    String res = allocReg(inst.getResult(), i);
                    asmList.add(String.format("\t%s %s, %s, %s\t\t#  %s", instName, res, lhs, rhs, inst));

                }
                case MOV -> {
                    String fmtStr;
                    String from;
                    // 先检查from的寄存器分配情况, 如果其为最后一次使用, res可以重复利用
                    if (inst.getFrom() instanceof IRImmediate imm) {
                        fmtStr = "\tli %s, %s\t\t#  %s";
                        from = imm.toString();
                    } else {
                        fmtStr = "\tmv %s, %s\t\t#  %s";
                        from = allocReg((IRVariable) inst.getFrom(), i);
                    }
                    String res = allocReg(inst.getResult(), i);
                    asmList.add(String.format(fmtStr, res, from, inst));
                }
                case RET -> {
                    if (inst.getReturnValue() instanceof IRImmediate imm) {
                        asmList.add(String.format("\tli a0, %s\t\t#  %s", imm, inst));
                    } else {
                        asmList.add(String.format("\tmv a0, %s\t\t#  %s", allocReg((IRVariable) inst.getReturnValue(), i), inst));
                    }
                }
            }
        }
    }


    /**
     * 输出汇编代码到文件
     *
     * @param path 输出文件路径
     */
    public void dump(String path) {
        // TODO: 输出汇编代码到文件
        FileUtils.writeLines(path, asmList);
    }
}

